package com.naja.auth2server.security;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.dao.DataRetrievalFailureException;
import org.springframework.lang.Nullable;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.OAuth2Token;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationCode;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @auther wangjianying
 * @date 2023/12/12 13:41
 */
public class RedisOAuth2AuthorizationService implements OAuth2AuthorizationService {
    private RegisteredClientRepository registeredClientRepository;
    private RedissonClient redisson;
    private ObjectMapper oauth2ObjectMapper;

    private final static Long STATE_TIMEOUT = 600L;

    private static final String AUTHORIZATION = "oauth2:token";

    public RedisOAuth2AuthorizationService(RedissonClient redisson, RegisteredClientRepository registeredClientRepository, ObjectMapper oauth2ObjectMapper) {
        this.oauth2ObjectMapper = oauth2ObjectMapper;
        this.redisson = redisson;
        this.registeredClientRepository = registeredClientRepository;
    }

    private Map<String,Object> getOAuth2AuthorizationRedisValue(OAuth2Authorization authorization) {
        Map<String,Object> parameters = new HashMap<>();
        parameters.put("id", authorization.getId());
        parameters.put("registeredClientId", authorization.getRegisteredClientId());
        parameters.put("principalName", authorization.getPrincipalName());
        parameters.put("authorizationGrantType", authorization.getAuthorizationGrantType().getValue());

        String authorizedScopes = null;
        if (!CollectionUtils.isEmpty(authorization.getAuthorizedScopes())) {
            authorizedScopes = StringUtils.collectionToDelimitedString(authorization.getAuthorizedScopes(), ",");
        }
        parameters.put("authorizedScopes", authorizedScopes);

        String attributes = writeMap(authorization.getAttributes());
        parameters.put("attributes", attributes);

        String state = null;
        String authorizationState = authorization.getAttribute(OAuth2ParameterNames.STATE);
        if (StringUtils.hasText(authorizationState)) {
            state = authorizationState;
        }
        parameters.put("state", state);

        OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode =
                authorization.getToken(OAuth2AuthorizationCode.class);
        Map<String, Object> authorizationCodeMap = toTokenMap("authorizationCode", authorizationCode);
        parameters.putAll(authorizationCodeMap);

        OAuth2Authorization.Token<OAuth2AccessToken> accessToken =
                authorization.getToken(OAuth2AccessToken.class);
        Map<String, Object> accessTokenMap = toTokenMap("accessToken" , accessToken);
        parameters.putAll(accessTokenMap);
        String accessTokenType = null;
        String accessTokenScopes = null;
        if (accessToken != null) {
            accessTokenType = accessToken.getToken().getTokenType().getValue();
            if (!CollectionUtils.isEmpty(accessToken.getToken().getScopes())) {
                accessTokenScopes = StringUtils.collectionToDelimitedString(accessToken.getToken().getScopes(), ",");
            }
        }
        parameters.put("accessTokenType", accessTokenType);
        parameters.put("accessTokenScopes", accessTokenScopes);

        OAuth2Authorization.Token<OidcIdToken> oidcIdToken = authorization.getToken(OidcIdToken.class);
        Map<String, Object> oidcIdTokenMap= toTokenMap("oidcIdToken", oidcIdToken);
        parameters.putAll(oidcIdTokenMap);

        OAuth2Authorization.Token<OAuth2RefreshToken> refreshToken = authorization.getRefreshToken();
        Map<String, Object> refreshTokenMap = toTokenMap("refreshToken", refreshToken);
        parameters.putAll(refreshTokenMap);
        return parameters;

    }

    private <T extends OAuth2Token> Map<String,Object> toTokenMap(String tokenColumnName, OAuth2Authorization.Token<T> token) {
        Map<String,Object> parameters = new HashMap<>();
        String tokenValue = null;
        Instant tokenIssuedAt = null;
        Instant tokenExpiresAt = null;
        String metadata = null;
        if (token != null) {
            tokenValue = token.getToken().getTokenValue();
            tokenIssuedAt = token.getToken().getIssuedAt();
            tokenExpiresAt = token.getToken().getExpiresAt();
            metadata = writeMap(token.getMetadata());
        }
        parameters.put(tokenColumnName+"Value", tokenValue);
        parameters.put(tokenColumnName+"IssuedAt", tokenIssuedAt);
        parameters.put(tokenColumnName+"ExpiresAt", tokenExpiresAt);
        parameters.put(tokenColumnName+"Metadata", metadata);
        return parameters;
    }

    private String writeMap(Map<String, Object> data) {
        try {
            return this.oauth2ObjectMapper.writeValueAsString(data);
        } catch (Exception ex) {
            throw new IllegalArgumentException(ex.getMessage(), ex);
        }
    }

    @Override
    public void save(OAuth2Authorization authorization) {
        Assert.notNull(authorization, "authorization cannot be null");
        Map<String, Object> oAuth2AuthorizationRedisValue = getOAuth2AuthorizationRedisValue(authorization);
        if (isState(authorization)) {
            String token = authorization.getAttribute("state");
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.STATE, token));
            bucket.set(oAuth2AuthorizationRedisValue,STATE_TIMEOUT, TimeUnit.SECONDS);
        }

        if (isCode(authorization)) {
            OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode = authorization.getToken(OAuth2AuthorizationCode.class);
            OAuth2AuthorizationCode authorizationCodeToken = authorizationCode.getToken();
            long between = ChronoUnit.SECONDS.between(authorizationCodeToken.getIssuedAt(), authorizationCodeToken.getExpiresAt());
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.CODE, authorizationCodeToken.getTokenValue()));
            bucket.set(oAuth2AuthorizationRedisValue,between, TimeUnit.SECONDS);
        }

        if (isRefreshToken(authorization)) {
            OAuth2RefreshToken refreshToken = authorization.getRefreshToken().getToken();
            long between = ChronoUnit.SECONDS.between(refreshToken.getIssuedAt(), refreshToken.getExpiresAt());
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.REFRESH_TOKEN, refreshToken.getTokenValue()));
            bucket.set(oAuth2AuthorizationRedisValue,between, TimeUnit.SECONDS);
        }

        if (isAccessToken(authorization)) {
            OAuth2AccessToken accessToken = authorization.getAccessToken().getToken();
            long between = ChronoUnit.SECONDS.between(accessToken.getIssuedAt(), accessToken.getExpiresAt());
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.ACCESS_TOKEN, accessToken.getTokenValue()));
            bucket.set(oAuth2AuthorizationRedisValue,between, TimeUnit.SECONDS);
        }
    }

    @Override
    public void remove(OAuth2Authorization authorization) {
        Assert.notNull(authorization, "authorization cannot be null");

        if (isState(authorization)) {
            String token = authorization.getAttribute("state");
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.STATE, token));
            bucket.deleteAsync();
        }

        if (isCode(authorization)) {
            OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode = authorization.getToken(OAuth2AuthorizationCode.class);
            OAuth2AuthorizationCode authorizationCodeToken = authorizationCode.getToken();
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.CODE, authorizationCodeToken.getTokenValue()));
            bucket.deleteAsync();
        }

        if (isRefreshToken(authorization)) {
            OAuth2RefreshToken refreshToken = authorization.getRefreshToken().getToken();
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.REFRESH_TOKEN, refreshToken.getTokenValue()));
            bucket.deleteAsync();
        }

        if (isAccessToken(authorization)) {
            OAuth2AccessToken accessToken = authorization.getAccessToken().getToken();
            RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(OAuth2ParameterNames.ACCESS_TOKEN, accessToken.getTokenValue()));
            bucket.deleteAsync();
        }
    }

    @Override
    @Nullable
    public OAuth2Authorization findById(String id) {
        throw new UnsupportedOperationException();
    }

    @Override
    @Nullable
    public OAuth2Authorization findByToken(String token, @Nullable OAuth2TokenType tokenType) {
        Assert.hasText(token, "token cannot be empty");
        Assert.notNull(tokenType, "tokenType cannot be empty");
        RBucket<Map<String, Object>> bucket = redisson.getBucket(buildKey(tokenType.getValue(), token));
        Map<String, Object> map = bucket.get();
        return mapRow(map);
    }

    public OAuth2Authorization mapRow(Map<String, Object> map) {
        String registeredClientId = (String) map.get("registeredClientId");
        RegisteredClient registeredClient = this.registeredClientRepository.findById(registeredClientId);
        if (registeredClient == null) {
            throw new DataRetrievalFailureException("The RegisteredClient with id '" + registeredClientId + "' was not found in the RegisteredClientRepository.");
        }
        OAuth2Authorization.Builder builder = OAuth2Authorization.withRegisteredClient(registeredClient);
        String id = (String) map.get("id");
        String principalName = (String) map.get("principalName");
        String authorizationGrantType = (String) map.get("authorizationGrantType");
        Set<String> authorizedScopes = Collections.emptySet();
        String authorizedScopesString = (String) map.get("authorizedScopes");
        if (authorizedScopesString != null) {
            authorizedScopes = StringUtils.commaDelimitedListToSet(authorizedScopesString);
        }
        Map<String, Object> attributes = this.parseMap((String) map.get("attributes"));
        builder.id(id)
                .principalName(principalName)
                .authorizationGrantType(new AuthorizationGrantType(authorizationGrantType))
                .authorizedScopes(authorizedScopes)
                .attributes((attrs) -> attrs.putAll(attributes));

        String state = (String) map.get("state");
        if (StringUtils.hasText(state)) {
            builder.attribute(OAuth2ParameterNames.STATE, state);
        }

        Instant tokenIssuedAt;
        Instant tokenExpiresAt;
        String authorizationCodeValue = (String) map.get("authorizationCodeValue");

        if (StringUtils.hasText(authorizationCodeValue)) {
            tokenIssuedAt = (Instant) map.get("authorizationCodeIssuedAt");
            tokenExpiresAt = (Instant) map.get("authorizationCodeExpiresAt");
            Map<String, Object> authorizationCodeMetadata = parseMap((String) map.get("authorizationCodeMetadata"));

            OAuth2AuthorizationCode authorizationCode = new OAuth2AuthorizationCode(
                    authorizationCodeValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(authorizationCode, (metadata) -> metadata.putAll(authorizationCodeMetadata));
        }

        String accessTokenValue = (String) map.get("accessTokenValue");
        if (StringUtils.hasText(accessTokenValue)) {
            tokenIssuedAt = (Instant) map.get("accessTokenIssuedAt");
            tokenExpiresAt = (Instant) map.get("accessTokenExpiresAt");
            Map<String, Object> accessTokenMetadata = parseMap((String) map.get("accessTokenMetadata"));
            OAuth2AccessToken.TokenType tokenType = null;
            if (OAuth2AccessToken.TokenType.BEARER.getValue().equalsIgnoreCase((String) map.get("accessTokenType"))) {
                tokenType = OAuth2AccessToken.TokenType.BEARER;
            }

            Set<String> scopes = Collections.emptySet();
            String accessTokenScopes = (String) map.get("accessTokenScopes");
            if (accessTokenScopes != null) {
                scopes = StringUtils.commaDelimitedListToSet(accessTokenScopes);
            }
            OAuth2AccessToken accessToken = new OAuth2AccessToken(tokenType, accessTokenValue, tokenIssuedAt, tokenExpiresAt, scopes);
            builder.token(accessToken, (metadata) -> metadata.putAll(accessTokenMetadata));
        }

        String oidcIdTokenValue = (String) map.get("oidcIdTokenValue");
        if (StringUtils.hasText(oidcIdTokenValue)) {
            tokenIssuedAt = (Instant) map.get("oidcIdTokenIssuedAt");
            tokenExpiresAt = (Instant) map.get("oidcIdTokenExpiresAt");
            Map<String, Object> oidcTokenMetadata = parseMap((String) map.get("oidcIdTokenMetadata"));

            OidcIdToken oidcToken = new OidcIdToken(
                    oidcIdTokenValue, tokenIssuedAt, tokenExpiresAt, (Map<String, Object>) oidcTokenMetadata.get(OAuth2Authorization.Token.CLAIMS_METADATA_NAME));
            builder.token(oidcToken, (metadata) -> metadata.putAll(oidcTokenMetadata));
        }

        String refreshTokenValue = (String) map.get("refreshTokenValue");
        if (StringUtils.hasText(refreshTokenValue)) {
            tokenIssuedAt = (Instant) map.get("refreshTokenIssuedAt");
            tokenExpiresAt = (Instant) map.get("refreshTokenExpiresAt");
            Map<String, Object> refreshTokenMetadata = parseMap((String) map.get("refreshTokenMetadata"));

            OAuth2RefreshToken refreshToken = new OAuth2RefreshToken(
                    refreshTokenValue, tokenIssuedAt, tokenExpiresAt);
            builder.token(refreshToken, (metadata) -> metadata.putAll(refreshTokenMetadata));
        }
        return builder.build();
    }
    private Map<String, Object> parseMap(String data) {
        try {
            return this.oauth2ObjectMapper.readValue(data, new TypeReference<Map<String, Object>>() {});
        } catch (Exception ex) {
            throw new IllegalArgumentException(ex.getMessage(), ex);
        }
    }

    private String buildKey(String type, String id) {
        return String.format("%s:%s:%s", AUTHORIZATION, type, id);
    }

    private static boolean isState(OAuth2Authorization authorization) {
        return Objects.nonNull(authorization.getAttribute("state"));
    }

    private static boolean isCode(OAuth2Authorization authorization) {
        OAuth2Authorization.Token<OAuth2AuthorizationCode> authorizationCode = authorization.getToken(OAuth2AuthorizationCode.class);
        return Objects.nonNull(authorizationCode);
    }

    private static boolean isRefreshToken(OAuth2Authorization authorization) {
        return Objects.nonNull(authorization.getRefreshToken());
    }

    private static boolean isAccessToken(OAuth2Authorization authorization) {
        return Objects.nonNull(authorization.getAccessToken());
    }


}
